from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, MessagesState, START

# 构建deepseek大模型客户端
llm = ChatOpenAI(
    model="deepseek-chat",
    base_url="https://api.deepseek.com/v1",
    openai_api_key="sk-ddad02a64c4c47ea9f6f526ef47cb602",
)


def call_model(state: MessagesState):
    """
    调用大语言模型处理消息状态

    Args:
        state (MessagesState): 包含对话消息的状态对象，其中messages字段存储着对话历史

    Returns:
        dict: 返回包含模型响应消息的字典，格式为{"messages": response}
    """
    response = llm.invoke(state["messages"])
    return {"messages": response}


# 创建状态图构建器，使用MessagesState作为状态类型
builder = StateGraph(MessagesState)

# 添加节点到状态图中
builder.add_node(call_model)
# 添加从起始节点到call_model节点的边
builder.add_edge(START, "call_model")

# 编译状态图生成可执行的图对象
graph = builder.compile()



#----------------------------------------------------
# for chunk in graph.stream(
#         {"messages": [{"role": "user", "content": "湖南的省会是哪⾥？"}]},
#         stream_mode="messages",
# ):
#     # print(chunk)
#     print(chunk[0].content)

#----------------------------------------------------

# 遍历图模型的流式响应，处理并打印消息内容
# 该循环通过graph.stream方法发送用户查询并接收流式响应
# 参数:
#   graph.stream() - 图模型的流式接口方法
#   {"messages": [{"role": "user", "content": "湖南的省会是哪⾥？"}]} - 包含用户消息的输入数据
#   stream_mode="messages" - 指定流式模式为消息模式
# 返回值:
#   message_chunk - 包含响应内容的消息块
#   metadata - 与消息相关的元数据
for message_chunk, metadata in graph.stream(
        {"messages": [{"role": "user", "content": "湖南的省会是哪⾥？"}]},
        stream_mode="messages",
):
    # 检查消息块是否具有content属性且内容不为空，然后打印内容
    if hasattr(message_chunk, 'content') and message_chunk.content:
        # 实时打印消息内容，不换行且立即刷新输出缓冲区
        print(message_chunk.content, end="", flush=True)
